from datasets import load_dataset, Audio
import os
import logging
from datetime import datetime
os.environ["HF_ENDPOINT"] = "https://hf-mirror.com"
#from datasets import load_dataset, Dataset, Audio
import os
# LOG_FILE = "./out_log/voxpopuli_processing.log"
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
LOG_FILE = f"./out_log/voxpopuli_processing_{timestamp}.log"
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
file_handler = logging.FileHandler(LOG_FILE, mode='w', encoding='utf-8')
console_handler = logging.StreamHandler()
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
file_handler.setFormatter(formatter)
console_handler.setFormatter(formatter)
logger.addHandler(file_handler)
logger.addHandler(console_handler)
logger.info("脚本开始执行。")
logger.info(f"日志将保存到文件: {LOG_FILE}")
# --- 步骤 1: 加载原始的 TSV 文件 ---
# 假设你的 TSV 文件就在 nl 文件夹里
tsv_file_path = "./voxpopuli-data/nl/asr_train.tsv"  # <-- 请替换成你的实际 TSV 文件名
logger.info(f"正在从 {tsv_file_path} 加载原始 TSV 文件...")
# 使用通用的 'csv' 加载器，并指定分隔符为 '\t'
raw_dataset = load_dataset('csv', data_files=tsv_file_path, delimiter='\t', split='train')
raw_dataset = raw_dataset.select(range(100))
logger.info(f"成功加载 {len(raw_dataset)} 个样本。")
# 检查初始加载的列名
logger.debug("原始 TSV 加载后的列名:")
logger.debug(raw_dataset.column_names)
# 预期输出: ['audio_id', 'raw_text', 'normalized_text', 'speaker_id', 'gender', 'is_gold_transcript', 'accent']

# --- 步骤 2: 添加 'language' 列 ---
# 为所有行添加一个固定值 "nl" 的新列
logger.info("正在为数据集添加 'language' 列...")
dataset_with_lang = raw_dataset.add_column('language', ['nl'] * len(raw_dataset))
logger.info("成功添加 'language' 列。")
logger.debug("添加 'language' 列后的列名:")
logger.debug(dataset_with_lang.column_names)

# --- 步骤 3: 构建 'audio' 列 ---
# 这是核心步骤。我们需要一个函数来根据 audio_id 生成完整的文件路径。

# 首先，定义你的音频文件所在的根目录
# 这个目录应该包含 "2014", "2015" 等年份子文件夹
audio_root = "./voxpopuli-data/nl/train/train_part_0/"  # <-- 请确认你的音频文件实际存放位置

def file_path_process():
    # 请将此路径修改为你的音频文件实际存放位置
    root_directory = "./voxpopuli-data/nl/train/"

    # 设置输出的 JSON 文件名
    output_json_file = "audio_files_index.json"

    # --- 2. 遍历目录并收集信息 ---
    # 创建一个空列表，用于存储所有文件的信息
    audio_files_data = {}

    logger.info(f"开始遍历音频目录: {audio_root} 以建立文件索引...")

    # os.walk 会生成目录树中的每个目录的元组 (dirpath, dirnames, filenames)
    for dirpath, dirnames, filenames in os.walk(root_directory):
        for filename in filenames:
            # 我们只关心 .wav 文件，可以根据需要修改扩展名
            if filename.lower().endswith(".wav"):
                # 获取文件的完整路径
                full_path = os.path.join(dirpath, filename)

                # 获取文件大小（字节）
                file_size = os.path.getsize(full_path)

                # 创建一个字典来存储当前文件的信息
                file_info = {
                    "filename": filename,
                    "path": full_path,
                    "size_bytes": file_size,
                    # 你可以在这里添加更多信息，比如从路径中提取的元数据
                    # "parent_folder": os.path.basename(dirpath)
                }

                # 将这个文件的信息字典添加到我们的列表中
                audio_files_data[filename.replace('.wav', '')] = file_info
    logger.info(f"遍历完成，共找到并索引了 {len(audio_files_data)} 个音频文件。")

    return audio_files_data
def build_audio_path(example):
    """
    这个函数接收数据集中的一行数据（一个 example 字典），
    并返回一个包含新 'audio' 列值的字典。
    """
    audio_id = example['id']



    # 使用 os.path.join 来构建路径，这能确保在不同操作系统下都有效
    full_path =file_path_process()[audio_id]

    # 返回一个新的字典，键是我们要创建的列名 'audio'，值是完整的路径
    return {'audio': full_path}


# 使用 map 函数将这个构建逻辑应用到数据集的每一行
# 这会创建一个名为 'audio' 的新列，其值为文件路径字符串
logger.info("正在使用 map 函数为每个样本构建 'audio' 路径...")
dataset_with_path = dataset_with_lang.map(build_audio_path)

logger.info("成功构建 'audio' 路径列。")
logger.debug("添加 'audio' 路径列后的列名:")
logger.debug(dataset_with_path.column_names)
logger.debug("'audio' 列的第一个元素（目前是路径字符串）:")
logger.debug(dataset_with_path[0]['audio'])
# 现在你可以看到 'audio' 列已经存在了，但它还是字符串类型


# --- 步骤 4: 将 'audio' 列转换为真正的 Audio 对象 ---
# 使用 cast_column 告诉 datasets 库，'audio' 列应该被当作音频文件来处理
logger.info("正在将 'audio' 列转换为 Audio 对象...")
final_dataset = dataset_with_path.cast_column('audio', Audio(sampling_rate=16000))

logger.info("成功转换 'audio' 列。")
logger.debug("'audio' 列的第一个元素（现在是 Audio 对象）:")
logger.debug(final_dataset[0]['audio'])
logger.debug(final_dataset.column_names)
# 列名没变，但列的类型已经变了


logger.info("正在加载 SpeechT5Processor...")
from transformers import SpeechT5Processor
checkpoint="microsoft/speecht5_tts"
processor=SpeechT5Processor.from_pretrained(checkpoint)
logger.info(f"成功从 '{checkpoint}' 加载 Processor。")
def extract_all_chars(batch):
    print("batch:",batch.keys())
    all_text=" ".join(batch["normalized_text"])
    vocab=list(set(all_text))
    return {"vocab": [vocab], "all_text":[all_text]}
#
vocabs=final_dataset.map(
    extract_all_chars,
    batched=True,
    batch_size=-1,
    keep_in_memory=True,
    remove_columns=final_dataset.column_names,
)
#
#
tokenizer=processor.tokenizer
dataset_vocab=set(vocabs["vocab"][0])
tokenizer_vocab={k for k, _ in tokenizer.get_vocab().items()}
logger.info("#### step1: 检查字符是否在 tokenizer 的词汇表中 ####")
logger.debug(f"Tokenizer 词汇表大小: {len(tokenizer_vocab)}")
logger.debug(f"数据集词汇表大小: {len(dataset_vocab)}")
unseen_chars = dataset_vocab - tokenizer_vocab
if unseen_chars:
    logger.warning(f"发现 {len(unseen_chars)} 个不在 tokenizer 词汇表中的字符: {unseen_chars}")
else:
    logger.info("数据集中的所有字符都存在于 tokenizer 的词汇表中。")


replacements=[
    ("à", "a"),
    ("ç", "c"),
    ("è", "e"),
    ("ë", "e"),
    ("í", "i"),
    ("ï", "i"),
    ("ö", "o"),
    ("ü", "u"),
]


def cleanup_text(inputs):
    for src, dst in replacements:
        inputs["normalized_text"]=inputs["normalized_text"].replace(src, dst)
        # print(inputs)
    return inputs
logger.info("#### step2: 非法字符校对及其数据清洗 #####")
final_dataset=final_dataset.map(cleanup_text)
logger.info("数据清洗完成。")
from collections import defaultdict
import matplotlib.pyplot as plt

speaker_counts=defaultdict(int)

for speaker_id in final_dataset["speaker_id"]:
    speaker_counts[speaker_id]+=1
logger.info("#### 检查数据集中的人员分布 ####")

logger.info(f"数据集中共有 {len(speaker_counts)} 位不同的说话人。")
logger.info("正在生成说话人样本数量分布图...")
#统计人员数量 便于构建人员语音特征 了解每个用户的语音构建贡献值 了解每个人的语音分布情况
plt.figure()
plt.hist(speaker_counts.values(), bins=20)
plt.ylabel("Speakers")
plt.xlabel("Examples")
plot_filename = "./out/speaker_distribution.png"
plt.savefig(plot_filename)
logger.info(f"说话人分布图已保存为 {plot_filename}")
plt.show()

#从语音数据集中筛选出“说话人样本数量适中”的数据，去除样本过少或过多的说话人，以实现数据集的平衡和训练效率的提升。
def select_speaker(speaker_id):
    return 100<=speaker_counts[speaker_id]<=400
#只有select_speaker为true才会返回数据保存数据
final_dataset_=final_dataset.filter(select_speaker, input_columns=["speaker_id"])

# Let's check how many speakers remain
# print(len(set(final_dataset["speaker_id"])))
logger.info(f"数据集中需要保留的数据有： {len(set(final_dataset['speaker_id']))} 位不同的说话人。")
# Let's see how many examples are left
logger.info(f"原本数据有： {len(final_dataset)} 条")
logger.info(f"剩余的数据有： {len(final_dataset_)} 条")  #只有0了 就先不过滤了
print(len(final_dataset))

#语音向量化 需要使用 SpeechBrain 提供的预训练模型 spkrec-xvect-voxceleb来生成向量化信息 包括每个人员的音色特征
#数据验证
logger.info("开始进行tts模型训练数据集构建")
import os
import torch
from speechbrain.pretrained import EncoderClassifier

spk_model_name = "speechbrain/spkrec-xvect-voxceleb"
out_model_path='./tmp'
if not os.path.exists(out_model_path):
    os.mkdir(out_model_path)
device = "cuda"
speaker_model = EncoderClassifier.from_hparams(
    source=spk_model_name,
    run_opts={"device": device},
    savedir=os.path.join(out_model_path, spk_model_name),
)
model_path = os.path.join(out_model_path, spk_model_name)
logger.info("语音向量化模型保存在：%s", model_path)

def create_speaker_embedding(waveform):
    '''
    输入：一段音频波形（waveform）
    输出：一个 512 维归一化的 speaker embedding
    用于后续控制 TTS 合成音色
    :param waveform:
    :return:
    '''
    with torch.no_grad():
        #关闭梯度计算 因为这里只是使用模型来进行向量化处理不需要更新参数进行反向传播
        speaker_embeddings = speaker_model.encode_batch(torch.tensor(waveform)) #将输入向量进行张量化后，进行编码压缩为512维
        speaker_embeddings = torch.nn.functional.normalize(speaker_embeddings, dim=2) #  F.normalize(x, p=2, dim=2) 这里p=2是默认值进行l2归一化处理 在第三维度上进行处理
        speaker_embeddings = speaker_embeddings.squeeze().cpu().numpy()

        # t_v=torch.tensor(waveform)
        # speaker_embeddings1 = speaker_model.encode_batch(t_v)
        # speaker_embeddings2 = torch.nn.functional.normalize(speaker_embeddings1, dim=2)
        # speaker_embeddings3 = speaker_embeddings2.squeeze()
        # speaker_embeddings4 = speaker_embeddings3.cpu()
        # speaker_embeddings5 = speaker_embeddings4.numpy()

        #squeeze去除维度为1的 转化为纯向量（1，1，512）——》（512，） .cpu()从gpu转到cpu
    return speaker_embeddings


def prepare_dataset(example):
    audio = example["audio"]

    example = processor(
        text=example["normalized_text"],
        audio_target=audio["array"],
        sampling_rate=audio["sampling_rate"],
        return_attention_mask=False,
    )

    # strip off thje batch dimension
    example["labels"] = example["labels"][0]

    # use SpeechBrain to obtain x-vector
    logger.info("正在开始speaker embedding...")
    example["speaker_embeddings"] = create_speaker_embedding(audio["array"])

    return example


processed_example=prepare_dataset(final_dataset[0])
logger.info(f"处理后的数据key值: {list(processed_example.keys())}")
logger.info(f"向量化的大小: {processed_example['speaker_embeddings'].shape}")
print('输出keys:',list(processed_example.keys()))

plt.figure()
plt.imshow(processed_example["labels"].T)
output_dir = "out"
os.makedirs(output_dir, exist_ok=True)
plot_path = os.path.join(output_dir, "mel_spectrogram.png")
plt.savefig(plot_path, dpi=300, bbox_inches="tight")
logger.info(f"梅尔频谱图已保存至: {plot_path}")
plt.show()

# # applying the processing function to the entire dataset
final_dataset=final_dataset.map(prepare_dataset, remove_columns=final_dataset.column_names)

#有些数据超长了 需要进行删除
def is_not_too_long(input_ids):
    input_length=len(input_ids)
    return input_length<200
logger.info("有些数据超长了 需要进行删除")
logger.info(f"超长数据处理,处理前{len(final_dataset)}")
final_dataset=final_dataset.filter(is_not_too_long, input_columns=["input_ids"])
logger.info(f"超长数据处理,处理后{len(final_dataset)}")

final_dataset=final_dataset.train_test_split(test_size=0.1)

logger.info("训练数据处理")
from dataclasses import dataclass
from typing import Any, Dict, List, Union


@dataclass
class TTSDataCollatorWithPadding:
    processor: Any
    '''
    我们希望将多个样本组合成一个批次（batch），因此需要定义一个自定义的数据整理器（custom data collator）。
    该整理器会使用填充符（padding tokens）对较短的序列进行填充，确保批次中所有样本具有相同的长度。
    对于频谱图标签（spectrogram labels），填充部分会被替换为特殊值 -100。这个特殊值会告诉模型在计算频谱图损失时忽略这些填充区域。
    将输入数据处理诚batch 批次处理
    输入：
    features = [
    {"input_ids": [...], "labels": [...], "speaker_embeddings": [...]},
    ...
    ]
    处理后的数据：
    batch = {
    "input_ids": (B, T_text),
    "labels": (B, T_spec, 80),
    "speaker_embeddings": (B, 512)
    }

    '''
    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
        input_ids = [{"input_ids": feature["input_ids"]} for feature in features]
        label_features = [{"input_values": feature["labels"]} for feature in features]
        speaker_features = [feature["speaker_embeddings"] for feature in features]
        # collate the inputs and targets into a batch
        batch = processor.pad(input_ids=input_ids, labels=label_features, return_tensors="pt")
        # replace padding with -100 to ignore loss correctly  将数据进行空缺数值的填补 将不需要的数据进行忽略 用-100进行处理
        batch["labels"] = batch["labels"].masked_fill(batch.decoder_attention_mask.unsqueeze(-1).ne(1), -100)
        # not used during fine-tuning
        del batch["decoder_attention_mask"]
        # round down target lengths to multiple of reduction factor
        if model.config.reduction_factor > 1:
            target_lengths = torch.tensor([len(feature["input_values"]) for feature in label_features])
            target_lengths = target_lengths.new(
                [length - length % model.config.reduction_factor for length in target_lengths]
            )
            #向下取整 避免维度错位
            max_length = max(target_lengths)
            batch["labels"] = batch["labels"][:, :max_length]

        # also add in the speaker embeddings
        batch["speaker_embeddings"] = torch.tensor(speaker_features)

        return batch
data_collator=TTSDataCollatorWithPadding(processor=processor)
from transformers import SpeechT5ForTextToSpeech
logger.info("训练模型加载")
model=SpeechT5ForTextToSpeech.from_pretrained(checkpoint)
model.config.use_cache=False #use_cache=true与模型训练时的梯度检查点不兼容 在模型训练时需要进行禁用 不开启
logger.info(f"模型的参数处理：{model.config}")
from transformers import Seq2SeqTrainingArguments, Seq2SeqTrainer

training_args=Seq2SeqTrainingArguments(
    output_dir="model_out",
    per_device_train_batch_size=4,
    gradient_accumulation_steps=8,
    learning_rate=1e-5,
    warmup_steps=50,
    max_steps=100,
    gradient_checkpointing=True,
    fp16=True,
    evaluation_strategy="steps",
    per_device_eval_batch_size=2,
    save_steps=50,
    eval_steps=50,
    logging_steps=25,
    report_to="none", # dryrun离线模式  report_to="wandb" or report_to="tensorboard" 我只需要在本地运行 不需要配置api来记录训练日志
    # report_to="wandb",
    run_name="model_out",
    load_best_model_at_end=True,
    greater_is_better=False,
    label_names=["labels"],
    push_to_hub=False,
)

trainer=Seq2SeqTrainer(
    args=training_args,
    model=model,
    train_dataset=final_dataset["train"],
    eval_dataset=final_dataset["test"],
    data_collator=data_collator,
    tokenizer=processor,
)
logger.info(f"模型开始训练")
trainer.train()
logger.info("模型推理")
# trainer.evaluate()
metrics = trainer.evaluate()
logger.info(f"评估指标:{metrics}")
# processor.push_to_hub(os.getenv("WANDB_NAME"))
# trainer.push_to_hub(os.getenv("WANDB_NAME"))

logger.info("模型推理")
#如果将模型上传至huggface可以使用以下推理
# from transformers import pipeline
# # pipe=pipeline("text-to-speech", model=os.getenv("WANDB_NAME"), device='cuda')
# pipe=pipeline("text-to-speech", model="./model_out/checkpoint-100", device='cuda')
# text="hallo allemaal, ik praat nederlands. groetjes aan iedereen!"
# example=final_dataset["test"][3]
# speaker_emebeddings=torch.tensor(example["speaker_emebddings"]).unsqueeze(0)
#
# forward_params={"speaker_embeddings": speaker_emebeddings}
# output=pipe(text, forward_params=forward_params)
# from IPython.display import Audio
# Audio(output['audio'], rate=output['sampling_rate'])

#本地推理
from transformers import SpeechT5ForTextToSpeech, SpeechT5Processor, SpeechT5HifiGan
import torch
from datasets import load_dataset
import soundfile as sf

# 加载模型和处理器
# model = SpeechT5ForTextToSpeech.from_pretrained("./model_out/checkpoint-100")
# processor = SpeechT5Processor.from_pretrained("./model_out/checkpoint-100")

model = SpeechT5ForTextToSpeech.from_pretrained("./model_out/checkpoint-100-0828")
processor = SpeechT5Processor.from_pretrained("./model_out/checkpoint-100-0828")
vocoder = SpeechT5HifiGan.from_pretrained("./tmp/speecht5_hifigan")

# 移动到GPU（如果可用）
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
vocoder.to(device)

# 加载示例数据获取说话人嵌入
example = final_dataset["test"][4]
print('example',example)
print("example.keys()",example.keys())
speaker_embeddings = torch.tensor(example["speaker_embeddings"]).unsqueeze(0).to(device)

# 处理文本输入
text = "hallo allemaal, ik praat nederlands. groetjes aan iedereen!"
inputs = processor(text=text, return_tensors="pt").to(device)

# 生成语音
with torch.no_grad():
    speech = model.generate_speech(
        inputs["input_ids"],
        speaker_embeddings,
        vocoder=vocoder
    )
# 将输出转换为音频数组
speech = speech.cpu().numpy()
# 保存音频到本地文件
output_filename = "./out/output_audio.wav"
sf.write(output_filename, speech, 16000)  # 确保使用正确的采样率
# 播放音频
from IPython.display import Audio
Audio(speech, rate=16000)  # 确保使用正确的采样率（可能与模型训练时相同）